import wfdb
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt

ECG_R_list = np.array(['N', 'f', 'e', '/', 'j', 'n', 'B',
                       'L', 'R', 'S', 'A', 'J', 'a', 'V',
                       'E', 'r', 'F', 'Q', '?'])


def extract_variable_length_beats(ecg_signal, r_peaks, symbols):
    """
    根据R波位置动态分割心拍。
    参数:
    ecg_signal: 一维ECG信号数组。
    r_peaks: R波峰值的索引列表。
    sampling_rate: 采样频率，默认为360Hz。
    返回:
    beats: 提取出的心拍列表，每个心拍是一个一维数组。
    beat_intervals: 每个心拍对应的起始和结束索引列表。
    """
    beats = []
    beat_intervals = []  # 存储每个心拍的起止点
    beat_symbol = []  # 存储每个心拍的标签

    # 计算RR间期
    rr_intervals = np.diff(r_peaks)

    for i in range(1, len(r_peaks) - 1):  # 跳过第一个和最后一个R波
        symbol = symbols[i]
        if symbol in ECG_R_list:
            rr_prev = rr_intervals[i - 1]  # 前一个RR间期
            rr_next = rr_intervals[i]  # 后一个RR间期
            # 动态计算心拍边界（使用40%/60%比例）
            start = r_peaks[i] - int(0.4 * rr_prev)
            end = r_peaks[i] + int(0.6 * rr_next)
            # 确保索引不超出信号范围
            if start < 0 or end > len(ecg_signal):
                print("报错了")
            beat = ecg_signal[start:end]
            beats.append(beat)
            beat_intervals.append((start, end))
            beat_symbol.append(symbol)
    return beats, beat_intervals, beat_symbol


# 五点平滑滤波
def fliter(ecg):
    # 五点平滑滤波
    ecg = np.convolve(ecg, np.ones((5,)) / 5, mode='same')
    # 工频干扰去除
    fs = 360  # 采样频率
    f0 = 50.0  # 要去除的工频干扰，50.0(抑制50Hz干扰)
    Q = 30.0  # 品质因数
    # 二阶IIR陷波数字滤波器，主要用于抑制信号中特定频率的干扰（如工频噪声）
    # w0，要抑制的中心频率（单位与 fs相同）；	品质因数，Q值越大，抑制带宽越窄； fs采样频率（必须提供，以明确频率单位）
    b, a = signal.iirnotch(w0=f0, Q=Q, fs=fs)
    # 工频干扰去除
    # 双向滤波函数，它通过正向和反向两次应用滤波器来实现零相位偏移，从而保持滤波后信号的相位特性不受影响
    ecg = signal.filtfilt(b, a, ecg)
    # 高通滤波器去除基线漂移
    [b, a] = signal.butter(3, [1 / 235 * 2, 20 / 235 * 2], 'bandpass')
    ecg = signal.filtfilt(b, a, ecg)
    return ecg


file_name = ['101']
for F_name in file_name:
    print(f'正在处理数据{F_name}')
    signal_annotation = wfdb.rdann(f'../data/mit-bih-arrhythmia-database-1.0.0/{F_name}', "atr", sampfrom=0,
                                   sampto=650000)
    record = wfdb.rdrecord(f'../data/mit-bih-arrhythmia-database-1.0.0/{F_name}', sampfrom=0, sampto=650000,
                           physical=True,
                           channels=[0])
    arr_transposed = record.p_signal.T
    arr1 = fliter(arr_transposed[0])
    print(arr1.shape)
    r_peaks = signal_annotation.sample
    variable_length_beats, intervals, beat_symbol = extract_variable_length_beats(arr1, r_peaks,
                                                                                  signal_annotation.symbol)
    # np.savetxt(f'../data/mit/1008/{F_name}-beta-data.csv', [variable_length_beats[0]], delimiter=',', fmt='%f')
    # np.savetxt(f'../data/mit/1008/{F_name}-beta-data.csv', [variable_length_beats[1]], delimiter=',', fmt='%f')
    # np.savetxt(f'../data/mit/1008/{F_name}-beta-data.csv', [variable_length_beats[2]], delimiter=',', fmt='%f')
    # np.savetxt(f'../data/mit/1008/{F_name}-beta-data.csv', [variable_length_beats[3]], delimiter=',', fmt='%f')
    # np.savetxt(f'../data/mit/1008/{F_name}-beta-data.csv', [variable_length_beats[4]], delimiter=',', fmt='%f')

    # 假设 variable_length_beats[0] 是您的心拍数据
    data = variable_length_beats[127]
    x_points = np.arange(len(data))  # 创建x轴坐标（采样点索引）

    # 方法1：使用彩色散点图（推荐）
    plt.figure(figsize=(12, 6))

    # 为每个点生成不同的颜色（使用viridis色彩映射）
    colors = plt.cm.viridis(np.linspace(0, 1, len(data)))

    plt.scatter(x_points, data, c=colors, s=40, alpha=0.8, edgecolors='black', linewidth=0.5)
    plt.plot(x_points, data, 'gray', alpha=0.3, linewidth=1)  # 添加灰色连线保持波形连续性

    plt.title('第一个等长心拍 - 多彩显示')
    plt.xlabel('采样点')
    plt.ylabel('振幅')
    plt.grid(True, alpha=0.3)
    plt.show()

    # # 方法2：基于振幅值的颜色映射
    # plt.figure(figsize=(12, 6))
    #
    # # 将振幅值归一化到0-1范围用于颜色映射
    # norm = plt.Normalize(vmin=np.min(data), vmax=np.max(data))
    # colors = plt.cm.plasma(norm(data))  # 使用plasma色彩映射
    #
    # scatter = plt.scatter(x_points, data, c=data, cmap='plasma', s=40, alpha=0.8)
    # plt.plot(x_points, data, 'lightgray', alpha=0.4, linewidth=1)  # 浅灰色背景线
    #
    # plt.colorbar(scatter, label='振幅值')  # 添加颜色条显示数值-颜色对应关系
    # plt.title('第一个等长心拍 - 振幅色彩映射')
    # plt.xlabel('采样点')
    # plt.ylabel('振幅')
    # plt.grid(True, alpha=0.2)
    # plt.show()